package com.wyt.spark.doris;

import org.apache.spark.sql.SparkSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.Statement;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;

/**
 * @Description:
 * @Author wangyongtao
 * @Date 2022/8/22 17:53
 **/
public class DorisUtils {

    private static Logger logger = LoggerFactory.getLogger(DorisUtils.class);

    public static SparkSession sparkSession;

    public static String url;
    public static Connection connection;
    public static Statement statement;

    public static boolean init() {
        try {
            String dorisUrl = "jdbc:mysql://172.172.178.76:9030/?rewriteBatchedStatements=true&allowMultiQueries=true";
            String username = "root";
            String password = "";
            Class.forName("com.mysql.cj.jdbc.Driver");
            connection = DriverManager.getConnection(dorisUrl, username, password);
            statement = connection.createStatement();
            url = dorisUrl;

        } catch (Exception e) {
            logger.error("初始化doris连接异常:", e);
            return false;
        }
        return true;
    }


    public static boolean truncateTable(String db, String table) {
        logger.info("doris删除全部数据的库表信息：{}.{}", db, table);
        String sql = "TRUNCATE TABLE " + db + "." + table;
        try {
            logger.info("doris删除全部数据sql为：{}", sql);
            statement.executeUpdate(sql);
        } catch (Exception e) {
            logger.error("doris删除全部数据异常：", e);
            return false;
        }
        return true;
    }

    public static boolean dropPartition(String db, String table, String startPartition) {
        logger.info("doris删除分区参数为：{}，{}，{}", db, table, startPartition);
        String partition_name = getPartiotionName(startPartition);
        String sql = "ALTER TABLE " + db + "." + table + " DROP PARTITION IF EXISTS " + partition_name;
        try {
            logger.info("doris删除分区sql为：{}", sql);
            statement.executeUpdate(sql);
        } catch (Exception e) {
            logger.error("doris删除分区[{}]异常，", startPartition, e);
            return false;
        }
        return true;
    }

    public static boolean createPartition(String db, String table, String startPartition, String stopPartition) {
        logger.info("doris创建分区参数为：{}，{}，{},{}", db, table, startPartition, stopPartition);
        String partition_name = "P" + startPartition.replaceAll("-", "");
        String sql = "ALTER TABLE " + db + "." + table + " ADD PARTITION IF NOT EXISTS " + partition_name + " VALUES [('" + startPartition + "'),('" + stopPartition + "'))";
        try {
            logger.info("doris创建分区sql为：{}", sql);
            statement.executeUpdate(sql);
        } catch (Exception e) {
            logger.error("doris创建分区[{}]异常，", partition_name, e);
            return false;
        }
        return true;
    }

    public static String buildInsertSql(String dbTable, String partition, String[] columns, Object[] values) {
        StringBuffer sb = new StringBuffer();
        sb.append("insert into ").append(dbTable).append(" PARTITION ( ").append(partition).append(")");

        sb.append(" (");
        for (String column : columns) {
            sb.append(column).append(",");
        }
        sb.deleteCharAt(sb.length() - 1);
        sb.append(")");

        sb.append(" VALUES (");
        for (Object value : values) {
            if (value instanceof String) {
                sb.append("'").append(value).append("'").append(",");
            } else {
                sb.append(value).append(",");
            }
        }
        sb.deleteCharAt(sb.length() - 1);
        sb.append(")");

        return sb.toString();
    }

    public static void insertBatch() throws Exception {
        String dbTable = "k_olap.test_doris_store";
        String[] columns = new String[]{"id", "name", "dt", "num", "year_str", "month_str", "day_str", "hour_str", "minute_str", "second_str", "precise"};
        Object[] values = new Object[11];

        LocalDateTime localDateTime = LocalDateTime.now();
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
        DateTimeFormatter partitionFormatter = DateTimeFormatter.ofPattern("yyyyMMdd");
        for (int i = 0; i < 1000000; i++) {
            String dateTimeStr = localDateTime.format(formatter);
            //构建参数
            values[0] = i + 1;
            values[1] = "name" + i;
            values[2] = dateTimeStr.substring(0, 10);
            values[3] = i;
            values[4] = dateTimeStr.substring(0, 4);
            values[5] = dateTimeStr.substring(5, 7);
            values[6] = dateTimeStr.substring(8, 10);
            values[7] = dateTimeStr.substring(11, 13);
            values[8] = dateTimeStr.substring(14, 16);
            values[9] = dateTimeStr.substring(17, 19);
            values[10] = i + 0.1;

            String sql = buildInsertSql(dbTable, "p" + localDateTime.format(partitionFormatter), columns, values);
            statement.execute(sql);

//            statement.addBatch(sql);
//            if (i > 0 && i % 500 == 0) {
//                statement.executeBatch();
//            }

            if (i > 0 && i % 10000 == 0) {
                //创建分区(已有初始分区)
                localDateTime = localDateTime.minusDays(1);
                createPartition("k_olap", "test_doris_store", localDateTime.format(formatter).substring(0, 10), dateTimeStr.substring(0, 10));
            }
        }
    }


    public static void close() {
        if (connection != null) {
            try {
                connection.close();
            } catch (Exception e) {
                logger.error("doris 连接关闭异常：", e);
            }
        }
        if (statement != null) {
            try {
                statement.close();
            } catch (Exception e) {
                logger.error("doris statement 关闭异常：", e);
            }
        }
    }

    private static String getPartiotionName(String startPartition) {
        String partition_name = "P" + startPartition.replaceAll("-", "");
        return partition_name;
    }


    public static void main(String[] args) throws Exception {
        init();
        insertBatch();
    }

}
